package cn.xiaopengstack.domain.agent.service.armory.node;

import cn.bugstack.wrench.design.framework.tree.StrategyHandler;
import cn.xiaopengstack.domain.agent.model.entity.ArmoryCommandEntity;
import cn.xiaopengstack.domain.agent.model.valobj.AiAgentEnumVO;
import cn.xiaopengstack.domain.agent.model.valobj.AiClientSystemPromptVO;
import cn.xiaopengstack.domain.agent.model.valobj.AiClientVO;
import cn.xiaopengstack.domain.agent.service.armory.AbstractArmorySupport;
import cn.xiaopengstack.domain.agent.service.armory.factory.DefaultArmoryStrategyFactory;
import com.alibaba.fastjson.JSONObject;
import io.modelcontextprotocol.client.McpSyncClient;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @author jiangyangang
 */
@Service
@Slf4j
public class AiClientNode extends AbstractArmorySupport {

    @Override
    protected String beanName(String id) {
        return AiAgentEnumVO.AI_CLIENT.getBeanName(id);
    }

    @Override
    protected String dataName() {
        return AiAgentEnumVO.AI_CLIENT.getDataName();
    }

    @Override
    protected String doApply(ArmoryCommandEntity requestParameter, DefaultArmoryStrategyFactory.DynamicContext dynamicContext) throws Exception {

        List<AiClientVO> clientVOList = dynamicContext.getValue(dataName());
        if (CollectionUtils.isEmpty(clientVOList)) {
            log.info("AI客户端加载，无可加载的AI Client...");
            return router(requestParameter, dynamicContext);
        }

        log.info("AI客户端加载，Client客户端加载中【{}】", JSONObject.toJSONString(requestParameter.getCommandIdList()));
        for (AiClientVO aiClientVO : clientVOList) {
            OpenAiChatModel openAiChatModel = getBean(aiClientVO.getModelBeanName());

            // 组装MCP
            List<String> mcpIdList = aiClientVO.getMcpBeanNameList();
            List<McpSyncClient> mcpSyncClients = mcpIdList.stream().map(beanName -> (McpSyncClient) getBean(beanName)).toList();

            // 组装系统提示词
            List<String> promptIdList = aiClientVO.getPromptIdList();
            Map<String, AiClientSystemPromptVO> systemPromptVOMap = dynamicContext.getValue(AiAgentEnumVO.AI_CLIENT_SYSTEM_PROMPT.getDataName());
            StringBuilder defaultSystem = new StringBuilder();
            for (String promptId : promptIdList) {
                AiClientSystemPromptVO systemPromptVO = systemPromptVOMap.getOrDefault(promptId, null);
                defaultSystem.append(Objects.isNull(systemPromptVO) ? "" : systemPromptVO.getPromptContent());
            }

            // 组装顾问
            List<String> advisorBeanNameList = aiClientVO.getAdvisorBeanNameList();
            List<Advisor> advisors = advisorBeanNameList.stream().map(beanName -> (Advisor) getBean(beanName)).toList();

            // 构建AI对话客户端
            ChatClient client = ChatClient.builder(openAiChatModel)
                    .defaultSystem(defaultSystem.toString())
                    .defaultAdvisors(advisors)
                    .defaultToolCallbacks(new SyncMcpToolCallbackProvider(mcpSyncClients).getToolCallbacks())
                    .build();
            registerBean(beanName(aiClientVO.getClientId()), ChatClient.class, client);
        }

        return router(requestParameter, dynamicContext);
    }

    @Override
    public StrategyHandler<ArmoryCommandEntity, DefaultArmoryStrategyFactory.DynamicContext, String> get(ArmoryCommandEntity requestParameter, DefaultArmoryStrategyFactory.DynamicContext dynamicContext) throws Exception {
        return defaultStrategyHandler;
    }
}
